# 1.导入Streamlit模块
import streamlit as st
from langchain_classic.chains import ConversationChain
from langchain_classic.memory import ConversationBufferMemory
from langchain_community.llms import Tongyi

# 1.先定义函数（功能接收用户问题，调用qwen-max模型，返回答案）
def get_chat_response(prompt,memory,apikey):
    # 获取模型对象
    tyqw=Tongyi(
        model_name="qwen-max",
        api_key=apikey
    )
    # 封装会话链
    chain = ConversationChain(llm=tyqw,memory=memory)
    # 发送请求，并获取答案
    res=chain.invoke({"input": prompt
    })
    # 返回答案
    return res["response"]

# 2.设置页面标题并添加分割线
st.title("通义聊天机器人")
# 添加分割线
st.divider()

# 3.先判断memory和messages是否在session_state中
# 这一块进行了修改，要分别判断是否在里面课程中只判断了memory是否在里面
if "memory" not in st.session_state:
    # 创建memory对象用于对话链
    st.session_state["memory"] = ConversationBufferMemory(return_messages=True)
if "messages" not in st.session_state:
    # 创建messages列表用于UI显示
    st.session_state["messages"] = [{'role':'ai','content':'你好，我是通义聊天机器人，有什么可以帮助你的么？'}]
# 4.聊天窗口
# （1）AI先说话
for message in st.session_state["messages"]:
    st.chat_message(message['role']).write(message['content'])
# （2）用户输入
prompt = st.chat_input()
# 添加左侧导航
with st.sidebar:
    st.title("导航栏")
    # 获取页面输入的api_key
    api_key = st.text_input("请输入您的通义API Key", type="password")
    st.markdown("""
    您可以在[通义控制台](https://bailian.console.aliyun.com/?apiKey=1#/api-key)获取您的API Key。
    """)
if prompt:
    # 如果用户没有输入api_key，就提示用户输入
    if not api_key:
        st.warning("请输入您的通义API Key")
        st.stop()
    # 展示用户提问的问题到页面，并添加到messages列表
    st.chat_message('user').write(prompt)
    st.session_state["messages"].append({'role':'user','content':prompt})
    with st.spinner("思考中..."):
        # 调用函数时传入正确的memory对象
        res=get_chat_response(prompt,st.session_state["memory"],api_key)
    # 展示AI的回答到页面，并添加到messages列表
    st.chat_message('ai').write(res)
    st.session_state["messages"].append({'role':'ai','content':res})

